
# 这个例子来自于：https://zhuanlan.zhihu.com/p/307349459

from scipy.sparse import dia_matrix
from scipy.sparse.linalg import inv
from numpy import pi
import numpy as np
import matplotlib.pyplot as plt  

class FEM:
    def __init__(self, nodes, xmin=0, xmax=1):
        self.nodes = nodes
        x = np.linspace(xmin, xmax, nodes)
        self.x = x
        self.h = x[1] - x[0]
        
    def Kmatrix(self):
        n = self.nodes
        m = 1/self.h * np.ones(n)
        data = [m, -2*m, m]
        offsets = [-1, 0, 1]
        # 使用 scipy.sparse 稀疏矩阵库，构造 3 对角稀疏矩阵 K
        K = dia_matrix((data, offsets), shape=(n, n)).tocsc()
        return K
    
    def bvec(self):
        '''Jackson 电动力学书本中介绍了一种简化方法，假设 $\rho(x)$在每个单元里是常数，只对$u_j(x)$积分，则积分结果为，'''
        return - np.sin(pi * self.x) * self.h
    
    def solve(self):
        K = self.Kmatrix()
        b = self.bvec()
        return inv(K) * b
    
    def compare(self):
        ground_truth = 1/(pi**2) * np.sin(pi * self.x)
        fem_res = self.solve()
        
        plt.plot(self.x, ground_truth, label="ground truth")
        plt.plot(self.x, fem_res, label="finite element")
        plt.title("number of nodes = %s"%self.nodes)
        plt.xlabel("x")
        plt.ylabel(r"$\phi(x)$")
        plt.legend(loc='best')
        plt.show()

fem = FEM(1001)
fem.compare()